package uima.stanford.parser;

import java.io.StringReader;
import java.util.ArrayList;
import java.util.List;

import org.apache.uima.analysis_component.JCasAnnotator_ImplBase;
import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.jcas.JCas;
import org.apache.uima.jcas.cas.FSArray;

import uima.stanford.types.UIMAConstituent;
import uima.stanford.types.UIMAToken;


import edu.stanford.nlp.ling.Word;
import edu.stanford.nlp.parser.lexparser.LexicalizedParser;
import edu.stanford.nlp.process.Tokenizer;
import edu.stanford.nlp.trees.GrammaticalStructure;
import edu.stanford.nlp.trees.GrammaticalStructureFactory;
import edu.stanford.nlp.trees.PennTreebankLanguagePack;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.trees.TreeGraphNode;
import edu.stanford.nlp.trees.TreePrint;

public class StanfordAE extends JCasAnnotator_ImplBase{
	
	private JCas jcas;
	private List<Word> wordList;
	
	@Override
	public void process(JCas arg0) throws AnalysisEngineProcessException {
		this.jcas = arg0;
		String grammarFilePath = (String) getContext().getConfigParameterValue("grammarFile");
		LexicalizedParser parser = new LexicalizedParser(grammarFilePath);
		
		String text = jcas.getDocumentText();
		parser.parse(text);
		//tokenizer
		PennTreebankLanguagePack tblp = new PennTreebankLanguagePack();
		Tokenizer<Word> tokenizer = tblp.getTokenizerFactory().getTokenizer(new StringReader(text));
		wordList = tokenizer.tokenize();		
		Tree parseTree =  (Tree) parser.apply(wordList);
		TreePrint treePrint =  new TreePrint("penn,wordsAndTags");
		//TreePrint treePrint =  new TreePrint("wordsAndTags");
		treePrint.printTree(parseTree);	
		
		GrammaticalStructureFactory gsf = tblp.grammaticalStructureFactory();
		GrammaticalStructure gs = gsf.newGrammaticalStructure(parseTree);	
		
		TreeGraphNode start = gs.root().children()[0];
		constructAnnotationTree(start);
	}
	
	private UIMAConstituent constructAnnotationTree(TreeGraphNode node){
		if(node.isPreTerminal()){			
			TreeGraphNode leaf = node.children()[0];
			int leafIndex = leaf.index();
			int begin = wordList.get(leafIndex-1).beginPosition();
			int end = wordList.get(leafIndex-1).endPosition();
			//construct uimatoken
			UIMAToken token = new UIMAToken(jcas);
			token.setBegin(begin);
			token.setEnd(end);
			token.setWord(leaf.label().value());			
			token.setTag(leaf.label().tag());
			//construct uimaconstituent
			UIMAConstituent constituent = new UIMAConstituent(jcas);
			constituent.setIsPos(true);
			constituent.setTag(node.label().value());
			constituent.setToken(token);
			constituent.setBegin(begin);
			constituent.setEnd(end);
			token.addToIndexes();
			constituent.addToIndexes();
			return constituent;			
		}
		else{
			UIMAConstituent constituent = new UIMAConstituent(jcas);
			constituent.setTag(node.label().value());
			constituent.setIsPos(false);
			TreeGraphNode[] children = node.children();
			int childrenNum = children.length;
			FSArray childrenArray = new FSArray(jcas, childrenNum);			
			for(int i = 0; i< childrenNum; i++){
				UIMAConstituent child = constructAnnotationTree(children[i]);
				//child.addToIndexes();
				childrenArray.set(i, child);				
			}
			constituent.setSubConstituents(childrenArray);
			ArrayList<TreeGraphNode> nodeLeaves = getLeaves(node);
			int firstIndex = nodeLeaves.get(0).index();
			int lastIndex = nodeLeaves.get(nodeLeaves.size()-1).index();
			constituent.setBegin(wordList.get(firstIndex-1).beginPosition());
			constituent.setEnd(wordList.get(lastIndex-1).endPosition());
			//constituent.addToIndexes();			
			return constituent;
		}
		
	}
	
	private ArrayList<TreeGraphNode> getLeaves(TreeGraphNode node){
		ArrayList<TreeGraphNode> leaves = new ArrayList<TreeGraphNode>();
		if(node.isPreTerminal()){
			
			if(node.isLeaf()){
				leaves.add(node);
			}
			else{
				leaves.add(node.children()[0]);
			}				
			return leaves;
		}
		else{
			TreeGraphNode[] children = node.children();
			ArrayList<TreeGraphNode> tmpLeavesFirst = getLeaves(children[0]);
			ArrayList<TreeGraphNode> tmpLeavesLast = getLeaves(children[children.length-1]);
			leaves.addAll(tmpLeavesFirst);
			leaves.addAll(tmpLeavesLast);
			return leaves;
		}
	}
	

}
